# Libraries Used
import os
import sys
import warnings

import numpy as np # Numerical computing
import pandas as pd # Data manipulation

import seaborn as sns # R-like high-level plots
import matplotlib.pyplot as plt # Python's base plotting 

from sklearn.decomposition import PCA

# home path
sys.path.append('../')
from src import utils, plots 
from src import tests, enrich

from ProteoForge import parse, process
from ProteoForge import impute, normalize

warnings.filterwarnings('ignore')
# Initialize the timer
startTime = utils.getTime()

## Configure Plotting Settings
def_colors = [
    "#139593", "#fca311", "#e54f2a",
    "#c3c3c3", "#555555",
    "#690000", "#5f4a00", "#004549"
]

condition_colors = {
    'True Hypoxia (1% O2) - 48 Hr': '#e9c46a',
    'True Hypoxia (1% O2) - 72 Hr': '#f4a261',
    'Normoxia (21% O2) - 48 Hr': '#ade8f4',
    'Normoxia (21% O2) - 72 Hr': '#0077b6',
    # Won't be used since the manuscript didn't cover these conditions
    # 'Chemical Hypoxia (CoCl2) - 48 Hr': '#c77dff',
    # 'Chemical Hypoxia (CoCl2) - 72 Hr': '#8338ec',
    # 'Oxidative Stress (H2O2) - 48 Hr': '#ff006e',
    # 'Oxidative Stress (H2O2) - 72 Hr': '#fb5607',
}

method_order = ['Original', 'mean(top)', 'mean(all)',  'dPFs']

proteinMethods = {
    "Original": "#f4f1de",
    "mean(top)": "#f2cc8f",
    "mean(all)": "#81b29a",
    "dPFs": "#183a37"
}

cv_group_palettes = {
    "<10%":  "#081c15",
    "<25%":  "#1b4332",
    "<50%": "#2d6a4f",
    "<100%": "#52b788",
    ">100%": "#95d5b2",
}

# --- Define Order and Colors ---
status_colors = {
    'Unexplained': "#C2C0C0",
    'Excluded': '#565d61',
    'Upregulated': '#780000',
    'Downregulated': '#e36414',
    'Equivalent': '#003049',
}

count_type_order = list(status_colors.keys())

status_markers = {
    "Equivalent": "o",
    "Upregulated": "^",
    "Downregulated": "v",
    "Unexplained": "s",
    "Excluded": "x"
}

cond_colors = {
    'Hypoxia (1% O2)': '#f4a261',
    'Normoxia (21% O2)': '#0077b6',
}
dura_colors = {
    '48hr': '#c3c3c3',
    '72hr': '#555555',
}
cond_order = ['Normoxia (21% O2)', 'Hypoxia (1% O2)']
dura_order = ['48hr', '72hr']

# Set seaborn style
sns.set_theme(
    style="white",
    context="paper",
    palette=def_colors,
    font_scale=1,
    rc={
        "figure.figsize": (6, 4),
        "font.family": "sans-serif",
        "font.sans-serif": ["Arial", "Ubuntu Mono"],
    }
)

# Figure Saving Settings
figure_formats = ["png", "pdf"]
save_to_folder = True
transparent_bg = True
figure_dpi = 300

## Configure dataframe displaying
pd.set_option('display.float_format', lambda x: '%.4f' % x)
pd.set_option('display.max_rows', 50)
pd.set_option('display.max_columns', 50)
pd.set_option('display.width', 500)  # Set a wider display width


## Printing Settings
verbose = True

plots.color_palette( def_colors, save=False )
plots.color_palette( condition_colors, save=False, name='condition', size=2)
plots.color_palette( proteinMethods, save=False, name='protein level calculation method', size=2)

home_path = './'
data_name = "hypoxia"
notebook_name = "04-Downstream"
data_path = f"{home_path}data/"
raw_path = f"{data_path}input/{data_name}/"
input_path = f"{data_path}cleaned/{data_name}/"
output_path = f"{data_path}results/{data_name}/"
figure_path = f"{home_path}figures/{data_name}/{notebook_name}/"

# Create the output folder if it does not exist
if not os.path.exists(output_path):
    os.makedirs(output_path)

# Create figure folder structure, if needed
if save_to_folder:
    for i in figure_formats:
        cur_folder = figure_path + i + "/"
        if not os.path.exists(cur_folder):
            os.makedirs(cur_folder)

peptide_df = pd.read_feather(f"{input_path}imputed_data.feather").reset_index()
proteins_used = peptide_df['Protein'].unique()

fasta_data = parse.fasta_to_df(
    fasta_path = f'./data/input/20230316_HumanCr_20806.fasta',
    gene_only = False,
    column_order=['entry', 'geneName', 'proteinDescription', 'sequenceLength', 'molecularWeight_kDa', 'sequence']
).rename(
    columns={
        'entry': 'Protein',
        'geneName': 'Gene',
        'proteinDescription': 'Description',
        'sequenceLength': 'Length',
        'molecularWeight_kDa': 'Weight(kDa)',
        'sequence': 'Sequence'
    }
)
fasta_data.head()

original_protein = pd.read_csv(f"{raw_path}report.pg_matrix.tsv", sep="\t", engine='pyarrow')
meta = original_protein.columns[4:].to_list()
# Build a meta dataframe
meta = pd.DataFrame(
    data=meta,
    columns=['filename']
)
# --- Clean and Standardize Metadata ---
print("\n🧹 Cleaning and standardizing metadata columns...")

meta['Run'] = meta['filename'].str.split(r'\\').str[-1].str.replace('.d', '')
# Extract condition identifier
def get_condition(x):
    if '1%' in x:
        return 'True Hypoxia (1% O2)'
    elif 'CoCl' in x:
        return 'Chemical Hypoxia (CoCl2)'
    elif 'H2O2' in x:
        return 'Oxidative Stress (H2O2)'
    else:
        return 'Normoxia (21% O2)'
meta['Condition'] = meta['Run'].map(get_condition)
# Extract the replicate number (special rule for 1% O2)
meta['Replicate'] = meta['Run'].str.split('_').apply(lambda x: x[2] if '1%' in x[1] else x[1][-1])
meta['Duration'] = meta['Run'].str.split('_').str[-5] # last 5th element from the end
meta['Colname'] = meta['Run'].str.split('_').map(
    # Simplify the colname based on the condition
    lambda x: f"{x[1]}_{x[3]}_{x[2]}" if '1%' in x[1] else f"{x[1][:-1]}_{x[2]}_{x[1][-1]}"
)
# --- Finalize DataFrame ---
print("Selecting final columns and renaming for clarity...")
print("\n✅ Metadata processing complete.")
print(f"Final DataFrame shape: {meta.shape}")
print(f"Final columns: {meta.columns.tolist()}")
print()

# Sort the metadata DataFrame for better readability
meta = meta.sort_values(
    ['Condition', 'Duration', 'Replicate'], 
    ascending=[False, True, True]
).reset_index(drop=True)

# Keep only the relevant conditions for the study
relevant_conditions = [
    'True Hypoxia (1% O2)',
    'Normoxia (21% O2)',
]
print(f"Filtering metadata to keep only relevant conditions: {relevant_conditions}")
meta = meta[meta['Condition'].isin(relevant_conditions)].reset_index(drop=True)
print(f"Filtered DataFrame shape: {meta.shape}")
meta

cv_datas = []
# Build a dictionary to match condition_colors keys: 'Condition - Duration Hr'
condition_to_samples = {}
for condition in meta['Condition'].unique():
    for duration in sorted(meta.loc[meta['Condition'] == condition, 'Duration'].unique()):
        key = f"{condition} - {duration} Hr"
        samples = meta.loc[(meta['Condition'] == condition) & (meta['Duration'] == duration), 'Colname'].tolist()
        condition_to_samples[key] = samples

# samples_to_condition mapping
samples_to_condition = {sample: condition for condition, samples in condition_to_samples.items() for sample in samples}
# samples_to_colors mapping
samples_to_colors = {sample: condition_colors[condition] for sample, condition in samples_to_condition.items()}

original_protein = pd.read_csv(f"{raw_path}report.pg_matrix.tsv", sep="\t", engine='pyarrow')
original_protein = original_protein.drop(
    columns=['Protein.Names', 'Genes', 'First.Protein.Description']
).rename(
    columns={"Protein.Group": "Protein"}
)
# Display summary information
print(f"Data shape (wide format): {original_protein.shape}")


print("\n--- 🧬 Protein ID Representation ---")
# Check if the protein contains '-' indicating it has isoforms
has_isoform = original_protein['Protein'].str.contains('-').any()
has_multiples = original_protein['Protein'].str.contains(';').any()
# Friendly print statements to explain the protein ID representation
if has_isoform:
    print("The dataset contains protein IDs with isoforms (e.g., 'P12345-2'). These represent specific variants of a protein.")
else:
    print("The dataset does not contain protein IDs with isoforms.")

if has_multiples:
    print("The dataset contains protein IDs with multiple entries (e.g., 'P12345; P67890'). These represent proteins that may have multiple forms or related proteins.")
    print()
    print("Will explode these entries to create a more detailed DataFrame.")
    original_protein = process.explode_aligned(original_protein, ['Protein'], sep=';', verbose=verbose)
    print()
    print(f"After exploding, the dataset now has {len(original_protein)} rows and {len(original_protein.columns)} columns.")
else:
    print("The dataset does not contain protein IDs with multiple entries.")

print(f"Values below 2 should be set as NaN, they might be artifacts of quantification.")
print(f"Number of missing values in the DataFrame before cleanup: {original_protein.isnull().sum().sum()}")
# Identify numeric columns in the DataFrame
numeric_columns = original_protein.select_dtypes(include=[np.number]).columns
original_protein[numeric_columns] = original_protein[numeric_columns].mask(
    original_protein[numeric_columns] < 2
)
print(f"Number of missing values in the DataFrame after cleanup: {original_protein.isnull().sum().sum()}")

original_protein = original_protein.set_index('Protein')[meta['filename']]
print("Set cleaned column names from metadata")
original_protein.columns = meta['Colname']
original_protein = original_protein.sort_index()
original_protein.head()

metric_weights = {  # Weights for the metrics used in pairwise combined score calculation
    'correlation': 1.0,
    'quantification': 1.0,
    'cv': 1.0
}
minOutlier_occurrance = 3 # minimum number of outlier occurrences to consider a sample behaves like an outlier

conditions_to_check = list(condition_to_samples.keys())  # List of conditions to check for outlier detection

print("\nCalculating Pairwise Combined Scores...")
print(" - Using metric weights:", metric_weights)
print(" - Using Spearman correlation, quantification as ratio, and CV as 1 - CV")

all_outlier_samples, combined_group_reports = utils.identify_outlier_samples(
    data=original_protein,
    sample_groups=condition_to_samples,
    analysis_func=utils.run_metric_combination_analysis,
    min_outlier_occurrence=minOutlier_occurrance,
    verbose=False
)
print(f"\nComplete list of unique outlier samples to remove (after min_outlier_occurrence >= {minOutlier_occurrance})\n - {all_outlier_samples}")
print()

plot_data = original_protein.reset_index().melt(
    id_vars='Protein',
    var_name='Sample',
    value_name='Intensity'
).dropna().assign(log2_Intensity=lambda x: np.log2(x['Intensity']))
fig, ax = plt.subplots(figsize=(12, 4))
sns.boxplot(
    x='Sample', y='log2_Intensity', hue='Sample', data=plot_data, ax=ax, 
    palette=samples_to_colors,
    fliersize=0.5,
    linewidth=0.5,
    notch=False,
    showmeans=True,
    meanprops={"marker": "o", "markerfacecolor": "white", "markeredgecolor": "black"},
    medianprops={"color": "white", "linewidth": 2.5},
    whiskerprops={"color": "black", "linewidth": 1},
    capprops={"color": "black", "linewidth": 0},
    legend=False
)
y_min = 0
y_max = plot_data['log2_Intensity'].max() + 1
ax.set_ylim(y_min, y_max)
ax.tick_params(axis='x', labelrotation=90)
ax.set_title('Boxplot of Log2 Intensities Across Samples', fontsize=16, fontweight='bold')
ax.set_xlabel('Sample', fontsize=14)
ax.set_ylabel('Log2(Intensity)', fontsize=14)
ax.grid(which='both', linestyle='--', linewidth=0.5, color='lightgray', alpha=0.7)

import matplotlib.patches as mpatches

# Highlight background for outlier samples
for i, label in enumerate(ax.get_xticklabels()):
    sample_label = label.get_text()
    if sample_label in all_outlier_samples:
        # Add a rectangle spanning the y-range at the x position
        ax.add_patch(
            mpatches.Rectangle(
                (i - 0.5, y_min),  # x position, y position
                1,                 # width (one boxplot)
                y_max - y_min,     # height
                color='#ffcccc',   # light red
                zorder=0,          # behind the boxplot
                alpha=0.5
            )
        )
# Add legend for outlier highlight
outlier_patch = mpatches.Patch(color='#ffcccc', label='Identified Outlier Sample', alpha=0.5)
ax.legend(handles=[outlier_patch], loc='upper right', fontsize=12)
plt.tight_layout()
plt.show()
# Remove the identified outlier samples from the main data
print("\nRemoving identified outlier samples from the main data...")
print(f"- Samples to remove: {all_outlier_samples}")
print(f"- Data shape before removal: {original_protein.shape}")
cleaned_data = original_protein.drop(columns=all_outlier_samples)
print(f"- Data shape after removal: {cleaned_data.shape}")
print("Outlier removal complete.")

meta = meta[~meta['Colname'].isin(all_outlier_samples)].reset_index(drop=True)
# Build a dictionary to match condition_colors keys: 'Condition - Duration Hr'
condition_to_samples = {}
for condition in meta['Condition'].unique():
    for duration in sorted(meta.loc[meta['Condition'] == condition, 'Duration'].unique()):
        key = f"{condition} - {duration} Hr"
        samples = meta.loc[(meta['Condition'] == condition) & (meta['Duration'] == duration), 'Colname'].tolist()
        condition_to_samples[key] = samples

# samples_to_condition mapping
samples_to_condition = {sample: condition for condition, samples in condition_to_samples.items() for sample in samples}
# samples_to_colors mapping
samples_to_colors = {sample: condition_colors[condition] for sample, condition in samples_to_condition.items()}

centered_data = normalize.by_median_centering(
    df=cleaned_data,
    rescale_to_original_magnitude=True,
    condition_map=None, # Pass dictionary if want to apply condition-specific centering
    verbose=verbose
)

imputation_pipeline_scheme = [
    { # First is the remove sparsely quantified features to set complete missing for low-value imputation 
        'method': 'amputate',
        'params': { 'min_quantified': 1 }
    },
    # # I generally don't use this step, but it can be useful for some datasets
    # { # Second step is generally to fill the sparse missingness with features' mean or median 
    #     'method': 'fill_dense', 
    #     'params': {'max_missing': 1, 'strategy': 'mean'}
    # },
    # This is the best for complete or mostly missing features
    { # Adds downshifted low-values to completely missing (1.0) or mostly missing (< 1.0) features
        'method': 'downshift',
        'params': {'missingness_threshold': .5, 'shift_magnitude': 2.5, 'low_percentile': 0.05}
    },
    { # Final step is to use k-NN imputation to fill the remaining missing values
        'method': 'knn',
        'params': {'n_neighbors': 4}
    }
]

imputed_data = impute.run_imputation_pipeline(
    data=centered_data,
    cond_dict=condition_to_samples,
    scheme=imputation_pipeline_scheme,
    is_log2=False,  
    return_log2=False,
    verbose=verbose
)
original_protein = imputed_data

print("\nCalculating Coefficient of Variation (CV) from proteins and group by condition...")

cv_data = pd.DataFrame(index = original_protein.index)
# Calculate the coefficient of variation per biological sample
for k, v in condition_to_samples.items():
    if not all(item in original_protein.columns for item in v):
        continue
    cv_data[k] = utils.cv_numpy(
        original_protein[v].values,
        axis=1,
        ignore_nan=True,
        format="percent"
    )

plot_data = utils.create_cv_group_plot_data(
    cv_data=cv_data,
    cv_group_palettes=cv_group_palettes
).dropna(subset=['CV'])
plot_data['Method'] = 'Original'
cv_datas.append(plot_data)

# PCA
scaler = "zscore"
complete_data = original_protein
scaled_data = utils.scale_the_data(
    complete_data,
    method=scaler,
    axis=1,
    is_log=False
).fillna(0)

# Run PCA
pca = PCA()
pca_res = pca.fit_transform(scaled_data.T)
pca_exp_var = pca.explained_variance_ratio_

# Set Variables
C1, C2 = 0, 1
hue_group = "Condition"
style_group = "Duration"
comp_locs = [C1, C2]
comp_names = [f'PC{c+1}' for c in comp_locs]
xlabel = f'{comp_names[0]} ({pca_exp_var[comp_locs[0]]:.2%} of variance)'
ylabel = f'{comp_names[1]} ({pca_exp_var[comp_locs[1]]:.2%} of variance)'

# Build the PCA Plot Data
plot_data = pd.DataFrame(pca_res[:, comp_locs], columns=comp_names)
plot_data['UniqueNames'] = complete_data.columns
plot_data['Condition'] = plot_data['UniqueNames'].map(samples_to_condition)
plot_data['Duration'] = plot_data['UniqueNames'].str.split('_').str[1]

# Find Centroids for each group (Condition and Duration)
centroid_data = plot_data.groupby([hue_group, style_group])[comp_names].mean().reset_index()

# Prepare color/marker mapping for unique (hue, style) combinations
unique_groups = centroid_data[[hue_group, style_group]].drop_duplicates()
group_palette = {}
group_marker = {}
markers = ["o", "s", "D", "^", "v", "P", "X", "*"]
for idx, (cond, dura) in enumerate(unique_groups.values):
    color = condition_colors.get(cond, "#888888")
    marker = markers[idx % len(markers)]
    group_palette[(cond, dura)] = color
    group_marker[(cond, dura)] = marker

fig, ax = plt.subplots(figsize=(9, 7))
fig.patch.set_facecolor('white')

# Plot technical replicates
for (cond, dura), group_df in plot_data.groupby([hue_group, style_group]):
    color = group_palette.get((cond, dura), "#888888")
    marker = group_marker.get((cond, dura), "o")
    ax.scatter(
        group_df[comp_names[0]], group_df[comp_names[1]],
        label=f"{cond}, {dura}hr",
        color=color, marker=marker, s=80, edgecolor='k', alpha=0.85, linewidth=1.2, zorder=2
    )
    # Confidence ellipse
    if len(group_df) > 1:
        plots.confidence_ellipse(
            group_df[comp_names[0]], group_df[comp_names[1]], ax, n_std=2,
            edgecolor=color, facecolor=color, alpha=0.18, linewidth=1.5, zorder=1
        )

# Plot centroids
for (cond, dura), group_df in centroid_data.groupby([hue_group, style_group]):
    color = group_palette.get((cond, dura), "#888888")
    marker = group_marker.get((cond, dura), "o")
    ax.scatter(
        group_df[comp_names[0]], group_df[comp_names[1]],
        color=color, marker=marker, s=260, edgecolor='k', linewidth=2.5, alpha=1, zorder=3
    )

# Axis limits and grid
minVal = plot_data[comp_names].min().min()
maxVal = plot_data[comp_names].max().max()
limit = max(abs(minVal), abs(maxVal))
offset = 0.1 * limit
limit += offset
ax.set_xlim(-limit, limit)
ax.set_ylim(-limit, limit)
ax.grid(True, which="both", linestyle="--", linewidth=0.7, alpha=0.4, color="grey")

# 0 lines
ax.axvline(0, color='grey', linestyle='--', linewidth=1.2, alpha=0.6, zorder=0)
ax.axhline(0, color='grey', linestyle='--', linewidth=1.2, alpha=0.6, zorder=0)

# Labels and title
ax.set_xlabel(xlabel, fontsize=14, fontweight="bold")
ax.set_ylabel(ylabel, fontsize=14, fontweight="bold")
ax.set_title(
    f"PCA of Original Protein Intensities (n={complete_data.shape[0]})",
    fontsize=17, fontweight="bold", pad=15, loc="left"
)

# Custom legend: combine color and marker
from matplotlib.lines import Line2D
legend_elements = [
    Line2D(
        [0], [0],
        marker=group_marker[(cond, dura)],
        color='w',
        markerfacecolor=group_palette[(cond, dura)],
        markeredgecolor='k',
        markersize=13,
        linewidth=0,
        label=f"{cond}, {dura}hr"
    )
    for (cond, dura) in unique_groups.values
]
ax.legend(
    handles=legend_elements,
    title="Condition, Duration",
    bbox_to_anchor=(1.02, 1),
    loc="upper left",
    frameon=False,
    fontsize=11,
    title_fontsize=12,
    borderaxespad=0.5
)

sns.despine(ax=ax, left=True, bottom=True)
plt.tight_layout()
plots.finalize_plot(
    fig, show=True, save=save_to_folder, 
    filename='PCA_Original_Proteins',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

peptide_df = pd.read_feather(f"{input_path}imputed_data.feather").reset_index()

# Ensure the first two columns are 'Protein' and 'unique_id', rest are samples
sample_cols = peptide_df.columns[2:]
# Calculate mean intensity for each peptide across all samples
peptide_df['mean_intensity'] = peptide_df[sample_cols].mean(axis=1, skipna=True)
# Sort by Protein and mean_intensity descending
peptide_df_sorted = peptide_df.sort_values(['Protein', 'mean_intensity'], ascending=[True, False])

# Select top 3 peptides per protein (vectorized)
top3_peptides = peptide_df_sorted.groupby('Protein', group_keys=False).head(3)

# Group by Protein and take the mean of sample columns (vectorized)
protein_top3 = top3_peptides.groupby('Protein')[sample_cols].mean()

print(f'Protein-level intensities (Top3 mean) shape: {protein_top3.shape}')

print(f"\n--- 🔄 Median Centering Normalization ---")
protein_top3 = normalize.by_median_centering(
    df=protein_top3,
    rescale_to_original_magnitude=True,
    condition_map=None,
    verbose=verbose
)

cv_data = pd.DataFrame(index = protein_top3.index)
# Calculate the coefficient of variation per biological sample
for k, v in condition_to_samples.items():
    # if not all(item in protein_top3.columns for item in v):
    #     continue
    cv_data[k] = utils.cv_numpy(
        protein_top3[v].values,
        axis=1,
        ignore_nan=True,
        format="percent"
    )

plot_data = utils.create_cv_group_plot_data(
    cv_data=cv_data,
    cv_group_palettes=cv_group_palettes
).dropna(subset=['CV'])
plot_data['Method'] = 'mean(top)'
cv_datas.append(plot_data)

protein_top3.head()

# PCA
scaler = "zscore"
complete_data = protein_top3.fillna(0)
scaled_data = utils.scale_the_data(
    complete_data,
    method=scaler,
    axis=1,
    is_log=False
)

# Run PCA
pca = PCA()
pca_res = pca.fit_transform(scaled_data.T)
pca_exp_var = pca.explained_variance_ratio_

# Set Variables
C1, C2 = 0, 1
hue_group = "Condition"
style_group = "Duration"
comp_locs = [C1, C2]
comp_names = [f'PC{c+1}' for c in comp_locs]
xlabel = f'{comp_names[0]} ({pca_exp_var[comp_locs[0]]:.2%} of variance)'
ylabel = f'{comp_names[1]} ({pca_exp_var[comp_locs[1]]:.2%} of variance)'

# Build the PCA Plot Data
plot_data = pd.DataFrame(pca_res[:, comp_locs], columns=comp_names)
plot_data['UniqueNames'] = complete_data.columns
plot_data['Condition'] = plot_data['UniqueNames'].map(samples_to_condition)
plot_data['Duration'] = plot_data['UniqueNames'].str.split('_').str[1]

# Find Centroids for each group (Condition and Duration)
centroid_data = plot_data.groupby([hue_group, style_group])[comp_names].mean().reset_index()

# Prepare color/marker mapping for unique (hue, style) combinations
unique_groups = centroid_data[[hue_group, style_group]].drop_duplicates()
group_palette = {}
group_marker = {}
markers = ["o", "s", "D", "^", "v", "P", "X", "*"]
for idx, (cond, dura) in enumerate(unique_groups.values):
    color = condition_colors.get(cond, "#888888")
    marker = markers[idx % len(markers)]
    group_palette[(cond, dura)] = color
    group_marker[(cond, dura)] = marker

fig, ax = plt.subplots(figsize=(9, 7))
fig.patch.set_facecolor('white')

# Plot technical replicates
for (cond, dura), group_df in plot_data.groupby([hue_group, style_group]):
    color = group_palette.get((cond, dura), "#888888")
    marker = group_marker.get((cond, dura), "o")
    ax.scatter(
        group_df[comp_names[0]], group_df[comp_names[1]],
        label=f"{cond}, {dura}hr",
        color=color, marker=marker, s=80, edgecolor='k', alpha=0.85, linewidth=1.2, zorder=2
    )
    # Confidence ellipse
    if len(group_df) > 1:
        plots.confidence_ellipse(
            group_df[comp_names[0]], group_df[comp_names[1]], ax, n_std=2,
            edgecolor=color, facecolor=color, alpha=0.18, linewidth=1.5, zorder=1
        )

# Plot centroids
for (cond, dura), group_df in centroid_data.groupby([hue_group, style_group]):
    color = group_palette.get((cond, dura), "#888888")
    marker = group_marker.get((cond, dura), "o")
    ax.scatter(
        group_df[comp_names[0]], group_df[comp_names[1]],
        color=color, marker=marker, s=260, edgecolor='k', linewidth=2.5, alpha=1, zorder=3
    )

# Axis limits and grid
minVal = plot_data[comp_names].min().min()
maxVal = plot_data[comp_names].max().max()
limit = max(abs(minVal), abs(maxVal))
offset = 0.1 * limit
limit += offset
ax.set_xlim(-limit, limit)
ax.set_ylim(-limit, limit)
ax.grid(True, which="both", linestyle="--", linewidth=0.7, alpha=0.4, color="grey")

# 0 lines
ax.axvline(0, color='grey', linestyle='--', linewidth=1.2, alpha=0.6, zorder=0)
ax.axhline(0, color='grey', linestyle='--', linewidth=1.2, alpha=0.6, zorder=0)

# Labels and title
ax.set_xlabel(xlabel, fontsize=14, fontweight="bold")
ax.set_ylabel(ylabel, fontsize=14, fontweight="bold")
ax.set_title(
    f"PCA of Averaging Top Peptides to Protein Intensities (n={complete_data.shape[0]})",
    fontsize=17, fontweight="bold", pad=15, loc="left"
)

# Custom legend: combine color and marker
from matplotlib.lines import Line2D
legend_elements = [
    Line2D(
        [0], [0],
        marker=group_marker[(cond, dura)],
        color='w',
        markerfacecolor=group_palette[(cond, dura)],
        markeredgecolor='k',
        markersize=13,
        linewidth=0,
        label=f"{cond}, {dura}hr"
    )
    for (cond, dura) in unique_groups.values
]
ax.legend(
    handles=legend_elements,
    title="Condition, Duration",
    bbox_to_anchor=(1.02, 1),
    loc="upper left",
    frameon=False,
    fontsize=11,
    title_fontsize=12,
    borderaxespad=0.5
)

sns.despine(ax=ax, left=True, bottom=True)
plt.tight_layout()
plots.finalize_plot(
    fig, show=True, save=save_to_folder, 
    filename='PCA_meanTop_Proteins',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

peptide_df = pd.read_feather(f"{input_path}imputed_data.feather").reset_index()
# Ensure the first two columns are 'Protein' and 'unique_id', rest are samples
sample_cols = peptide_df.columns[2:]

# Average all peptides per protein
protein_all = peptide_df.groupby('Protein')[sample_cols].mean()

print(f'Protein-level intensities (All peptides mean) shape: {protein_all.shape}')

print(f"\n--- 🔄 Median Centering Normalization ---")
protein_all = normalize.by_median_centering(
    df=protein_all,
    rescale_to_original_magnitude=True,
    condition_map=None,
    verbose=verbose
)

cv_data = pd.DataFrame(index = protein_all.index)
# Calculate the coefficient of variation per biological sample
for k, v in condition_to_samples.items():
    # if not all(item in protein_all.columns for item in v):
    #     continue
    cv_data[k] = utils.cv_numpy(
        protein_all[v].values,
        axis=1,
        ignore_nan=True,
        format="percent"
    )

plot_data = utils.create_cv_group_plot_data(
    cv_data=cv_data,
    cv_group_palettes=cv_group_palettes
).dropna(subset=['CV'])
plot_data['Method'] = 'mean(all)'
cv_datas.append(plot_data)


protein_all.head()

# PCA
scaler = "zscore"
complete_data = protein_all.fillna(0)
scaled_data = utils.scale_the_data(
    complete_data,
    method=scaler,
    axis=1,
    is_log=False
)

# Run PCA
pca = PCA()
pca_res = pca.fit_transform(scaled_data.T)
pca_exp_var = pca.explained_variance_ratio_

# Set Variables
C1, C2 = 0, 1
hue_group = "Condition"
style_group = "Duration"
comp_locs = [C1, C2]
comp_names = [f'PC{c+1}' for c in comp_locs]
xlabel = f'{comp_names[0]} ({pca_exp_var[comp_locs[0]]:.2%} of variance)'
ylabel = f'{comp_names[1]} ({pca_exp_var[comp_locs[1]]:.2%} of variance)'

# Build the PCA Plot Data
plot_data = pd.DataFrame(pca_res[:, comp_locs], columns=comp_names)
plot_data['UniqueNames'] = complete_data.columns
plot_data['Condition'] = plot_data['UniqueNames'].map(samples_to_condition)
plot_data['Duration'] = plot_data['UniqueNames'].str.split('_').str[1]

# Find Centroids for each group (Condition and Duration)
centroid_data = plot_data.groupby([hue_group, style_group])[comp_names].mean().reset_index()

# Prepare color/marker mapping for unique (hue, style) combinations
unique_groups = centroid_data[[hue_group, style_group]].drop_duplicates()
group_palette = {}
group_marker = {}
markers = ["o", "s", "D", "^", "v", "P", "X", "*"]
for idx, (cond, dura) in enumerate(unique_groups.values):
    color = condition_colors.get(cond, "#888888")
    marker = markers[idx % len(markers)]
    group_palette[(cond, dura)] = color
    group_marker[(cond, dura)] = marker

fig, ax = plt.subplots(figsize=(9, 7))
fig.patch.set_facecolor('white')

# Plot technical replicates
for (cond, dura), group_df in plot_data.groupby([hue_group, style_group]):
    color = group_palette.get((cond, dura), "#888888")
    marker = group_marker.get((cond, dura), "o")
    ax.scatter(
        group_df[comp_names[0]], group_df[comp_names[1]],
        label=f"{cond}, {dura}hr",
        color=color, marker=marker, s=80, edgecolor='k', alpha=0.85, linewidth=1.2, zorder=2
    )
    # Confidence ellipse
    if len(group_df) > 1:
        plots.confidence_ellipse(
            group_df[comp_names[0]], group_df[comp_names[1]], ax, n_std=2,
            edgecolor=color, facecolor=color, alpha=0.18, linewidth=1.5, zorder=1
        )

# Plot centroids
for (cond, dura), group_df in centroid_data.groupby([hue_group, style_group]):
    color = group_palette.get((cond, dura), "#888888")
    marker = group_marker.get((cond, dura), "o")
    ax.scatter(
        group_df[comp_names[0]], group_df[comp_names[1]],
        color=color, marker=marker, s=260, edgecolor='k', linewidth=2.5, alpha=1, zorder=3
    )

# Axis limits and grid
minVal = plot_data[comp_names].min().min()
maxVal = plot_data[comp_names].max().max()
limit = max(abs(minVal), abs(maxVal))
offset = 0.1 * limit
limit += offset
ax.set_xlim(-limit, limit)
ax.set_ylim(-limit, limit)
ax.grid(True, which="both", linestyle="--", linewidth=0.7, alpha=0.4, color="grey")

# 0 lines
ax.axvline(0, color='grey', linestyle='--', linewidth=1.2, alpha=0.6, zorder=0)
ax.axhline(0, color='grey', linestyle='--', linewidth=1.2, alpha=0.6, zorder=0)

# Labels and title
ax.set_xlabel(xlabel, fontsize=14, fontweight="bold")
ax.set_ylabel(ylabel, fontsize=14, fontweight="bold")
ax.set_title(
    f"PCA of Averaging All Peptides to Protein Intensities (n={complete_data.shape[0]})",
    fontsize=17, fontweight="bold", pad=15, loc="left"
)

# Custom legend: combine color and marker
from matplotlib.lines import Line2D
legend_elements = [
    Line2D(
        [0], [0],
        marker=group_marker[(cond, dura)],
        color='w',
        markerfacecolor=group_palette[(cond, dura)],
        markeredgecolor='k',
        markersize=13,
        linewidth=0,
        label=f"{cond}, {dura}hr"
    )
    for (cond, dura) in unique_groups.values
]
ax.legend(
    handles=legend_elements,
    title="Condition, Duration",
    bbox_to_anchor=(1.02, 1),
    loc="upper left",
    frameon=False,
    fontsize=11,
    title_fontsize=12,
    borderaxespad=0.5
)

sns.despine(ax=ax, left=True, bottom=True)
plt.tight_layout()
plots.finalize_plot(
    fig, show=True, save=save_to_folder, 
    filename='PCA_meanAll_Proteins',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

dPF_long = pd.read_feather(f"{output_path}test_data_48hr.feather")

# Remove the PTM peptides (dPF=-1) since they don't contribute to the protein quantification
dPF_long = dPF_long[dPF_long['dPF'] != -1].reset_index(drop=True)
dPF_long['Proteoform'] = dPF_long['Protein'] + "_" + dPF_long['dPF'].astype(str)
protein_dpf_48 = dPF_long.pivot_table(
    index='Proteoform',
    columns='Sample',
    values='Intensity',
    aggfunc='mean'  # In case of duplicates, take the mean
)

print(f'Protein-level intensities (Proteoform dPF) shape: {protein_dpf_48.shape}')
print(f"\n--- 🔄 Median Centering Normalization ---")
protein_dpf_48 = normalize.by_median_centering(
    df=protein_dpf_48,
    rescale_to_original_magnitude=True,
    # condition_map={k: v for k, v in condition_to_samples.items() if '48 Hr' in k}, 
    condition_map=None,
    verbose=verbose
)



cv_data = pd.DataFrame(index = protein_dpf_48.index)
# Calculate the coefficient of variation per biological sample
for k, v in condition_to_samples.items():
    if not all(item in protein_dpf_48.columns for item in v):
        continue
    cv_data[k] = utils.cv_numpy(
        protein_dpf_48[v].values,
        axis=1,
        ignore_nan=True,
        format="percent"
    )

plot_data = utils.create_cv_group_plot_data(
    cv_data=cv_data,
    cv_group_palettes=cv_group_palettes,
    id_col='Proteoform'
).dropna(subset=['CV'])
plot_data['Method'] = 'dPFs'
cv_datas.append(plot_data)

dPF_long = pd.read_feather(f"{output_path}test_data_72hr.feather")

# Remove the PTM peptides (dPF=-1) since they don't contribute to the protein quantification
dPF_long = dPF_long[dPF_long['dPF'] != -1].reset_index(drop=True)
dPF_long['Proteoform'] = dPF_long['Protein'] + "_" + dPF_long['dPF'].astype(str)
protein_dpf_72 = dPF_long.pivot_table(
    index='Proteoform',
    columns='Sample',
    values='Intensity',
    aggfunc='mean'  # In case of duplicates, take the mean
)

print(f'Protein-level intensities (Proteoform dPF) shape: {protein_dpf_72.shape}')
print(f"\n--- 🔄 Median Centering Normalization ---")
protein_dpf_72 = normalize.by_median_centering(
    df=protein_dpf_72,
    rescale_to_original_magnitude=True,
    # condition_map={k: v for k, v in condition_to_samples.items() if '72 Hr' in k}, 
    condition_map=None,
    verbose=verbose
)

cv_data = pd.DataFrame(index = protein_dpf_72.index)
# Calculate the coefficient of variation per biological sample
for k, v in condition_to_samples.items():
    if not all(item in protein_dpf_72.columns for item in v):
        continue
    cv_data[k] = utils.cv_numpy(
        protein_dpf_72[v].values,
        axis=1,
        ignore_nan=True,
        format="percent"
    )

plot_data = utils.create_cv_group_plot_data(
    cv_data=cv_data,
    cv_group_palettes=cv_group_palettes,
    id_col='Proteoform'
).dropna(subset=['CV'])

plot_data['Method'] = 'dPFs'
cv_datas.append(plot_data)

dPF_Proteins = pd.concat([protein_dpf_48, protein_dpf_72], axis=1)
# PCA
scaler = "zscore"
complete_data = dPF_Proteins.fillna(0)
scaled_data = utils.scale_the_data(
    complete_data,
    method=scaler,
    axis=1,
    is_log=False
)

# Run PCA
pca = PCA()
pca_res = pca.fit_transform(scaled_data.T)
pca_exp_var = pca.explained_variance_ratio_

# Set Variables
C1, C2 = 0, 1
hue_group = "Condition"
style_group = "Duration"
comp_locs = [C1, C2]
comp_names = [f'PC{c+1}' for c in comp_locs]
xlabel = f'{comp_names[0]} ({pca_exp_var[comp_locs[0]]:.2%} of variance)'
ylabel = f'{comp_names[1]} ({pca_exp_var[comp_locs[1]]:.2%} of variance)'

# Build the PCA Plot Data
plot_data = pd.DataFrame(pca_res[:, comp_locs], columns=comp_names)
plot_data['UniqueNames'] = complete_data.columns
plot_data['Condition'] = plot_data['UniqueNames'].map(samples_to_condition)
plot_data['Duration'] = plot_data['UniqueNames'].str.split('_').str[1]

# Find Centroids for each group (Condition and Duration)
centroid_data = plot_data.groupby([hue_group, style_group])[comp_names].mean().reset_index()

# Prepare color/marker mapping for unique (hue, style) combinations
unique_groups = centroid_data[[hue_group, style_group]].drop_duplicates()
group_palette = {}
group_marker = {}
markers = ["o", "s", "D", "^", "v", "P", "X", "*"]
for idx, (cond, dura) in enumerate(unique_groups.values):
    color = condition_colors.get(cond, "#888888")
    marker = markers[idx % len(markers)]
    group_palette[(cond, dura)] = color
    group_marker[(cond, dura)] = marker

fig, ax = plt.subplots(figsize=(9, 7))
fig.patch.set_facecolor('white')

# Plot technical replicates
for (cond, dura), group_df in plot_data.groupby([hue_group, style_group]):
    color = group_palette.get((cond, dura), "#888888")
    marker = group_marker.get((cond, dura), "o")
    ax.scatter(
        group_df[comp_names[0]], group_df[comp_names[1]],
        label=f"{cond}, {dura}hr",
        color=color, marker=marker, s=80, edgecolor='k', alpha=0.85, linewidth=1.2, zorder=2
    )
    # Confidence ellipse
    if len(group_df) > 1:
        plots.confidence_ellipse(
            group_df[comp_names[0]], group_df[comp_names[1]], ax, n_std=2.5,
            edgecolor=color, facecolor=color, alpha=0.18, linewidth=1.5, zorder=1
        )

# Plot centroids
for (cond, dura), group_df in centroid_data.groupby([hue_group, style_group]):
    color = group_palette.get((cond, dura), "#888888")
    marker = group_marker.get((cond, dura), "o")
    ax.scatter(
        group_df[comp_names[0]], group_df[comp_names[1]],
        color=color, marker=marker, s=260, edgecolor='k', linewidth=2.5, alpha=1, zorder=3
    )

# Axis limits and grid
minVal = plot_data[comp_names].min().min()
maxVal = plot_data[comp_names].max().max()
limit = max(abs(minVal), abs(maxVal))
# offset = 0.10 * limit
# limit += offset
# ax.set_xlim(-limit, limit)
# ax.set_ylim(-limit, limit)
ax.grid(True, which="both", linestyle="--", linewidth=0.7, alpha=0.4, color="grey")

# 0 lines
ax.axvline(0, color='grey', linestyle='--', linewidth=1.2, alpha=0.6, zorder=0)
ax.axhline(0, color='grey', linestyle='--', linewidth=1.2, alpha=0.6, zorder=0)

# Labels and title
ax.set_xlabel(xlabel, fontsize=14, fontweight="bold")
ax.set_ylabel(ylabel, fontsize=14, fontweight="bold")
ax.set_title(
    f"PCA of Averaging dPF Peptides (n={complete_data.shape[0]})",
    fontsize=17, fontweight="bold", pad=15, loc="left"
)

# Custom legend: combine color and marker
from matplotlib.lines import Line2D
legend_elements = [
    Line2D(
        [0], [0],
        marker=group_marker[(cond, dura)],
        color='w',
        markerfacecolor=group_palette[(cond, dura)],
        markeredgecolor='k',
        markersize=13,
        linewidth=0,
        label=f"{cond}, {dura}hr"
    )
    for (cond, dura) in unique_groups.values
]
ax.legend(
    handles=legend_elements,
    title="Condition, Duration",
    bbox_to_anchor=(1.02, 1),
    loc="upper left",
    frameon=False,
    fontsize=11,
    title_fontsize=12,
    borderaxespad=0.5
)

sns.despine(ax=ax, left=True, bottom=True)
plt.tight_layout()
plots.finalize_plot(
    fig, show=True, save=save_to_folder, 
    filename='PCA_dPF_Proteins',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

variance_data = []
proteinDatasets = {
    "Original": original_protein,
    "mean(top)": protein_top3,
    "mean(all)": protein_all,
    "dPFs": dPF_Proteins
}
for k, v in proteinDatasets.items():
    # Perform PCA on the centered data
    
    # Drop any missing values
    complete_data = v
    # Scale the data
    scaled_data = utils.scale_the_data(
        complete_data,
        method='zscore',
        axis=1,
        is_log=False
    ).fillna(0)
    # Run the PCA
    pca = PCA()
    pca_res = pca.fit_transform(scaled_data.T)
    pca_exp_var = pca.explained_variance_ratio_
    pca_exp_var_cum = np.cumsum(pca_exp_var)
    for i, cum_var in enumerate(pca_exp_var_cum):
        variance_data.append({
            "Data": k,
            "PC Component": i + 1,
            "Cumulative Variance": cum_var
        })

plot_data = pd.DataFrame(variance_data)

# Improved cumulative explained variance plot for PCA

import matplotlib.ticker as mtick

fig, ax = plt.subplots(figsize=(7, 6))

# Robust lineplot with markers and edgecolor for visibility
sns.lineplot(
    data=plot_data,
    x="PC Component",
    y="Cumulative Variance",
    hue="Data",
    style="Data",
    palette=proteinMethods,
    dashes=False,
    linewidth=2.5,
    markers=True,
    markersize=9,
    alpha=0.85,
    ax=ax,
    legend="full",
    markeredgecolor="grey"  # Add edgecolor for marker visibility
)

# Add horizontal reference lines for key variance thresholds
for thresh, color in zip([0.8, 0.9, 0.95], ['#e9c46a', '#f4a261', '#0077b6']):
    ax.axhline(thresh, color=color, linestyle='--', linewidth=1.2, alpha=0.5, zorder=0)
    ax.text(
        plot_data["PC Component"].max() + 0.2, thresh, f"{int(thresh*100)}%",
        va='center', ha='left', color=color, fontsize=11, fontweight='bold', alpha=0.7
    )

# Annotate the number of PCs needed to reach 80% and 90% variance for each method
for method in plot_data['Data'].unique():
    method_data = plot_data[plot_data['Data'] == method]
    for thresh in [0.8, 0.9]:
        pcs = method_data[method_data['Cumulative Variance'] >= thresh]['PC Component']
        if not pcs.empty:
            pc_num = pcs.iloc[0]
            var_val = method_data[method_data['PC Component'] == pc_num]['Cumulative Variance'].values[0]
            ax.scatter(pc_num, var_val, color=proteinMethods[method], s=80, edgecolor='k', zorder=5)
            ax.text(
                pc_num, var_val + 0.03, f"{pc_num} PCs", 
                ha='center', va='bottom', fontsize=10, fontweight='semibold', color=proteinMethods[method]
            )

# Set axis limits and ticks
ax.set_ylim(-0.05, 1.05)
ax.set_xlim(1, plot_data["PC Component"].max())
ax.set_xticks(range(1, plot_data["PC Component"].max() + 1))
ax.set_xticklabels(
    [str(i) for i in range(1, plot_data["PC Component"].max() + 1)],
    rotation=0, ha="center", fontsize=11, fontweight="bold"
)

# Format y-axis as percent
ax.yaxis.set_major_formatter(mtick.PercentFormatter(1.0, decimals=0))

# Axis labels and title
ax.set_xlabel("Principal Component Number", fontsize=14, fontweight="bold", labelpad=10)
ax.set_ylabel("Cumulative Explained Variance (%)", fontsize=14, fontweight="bold", labelpad=10)
ax.set_title("Cumulative Explained Variance by PCA\nfor Different Protein Quantification Methods", fontsize=16, fontweight="bold", pad=15)

# Grid and background
ax.grid(True, which="both", linestyle="--", linewidth=0.7, alpha=0.5, color="lightgrey")
ax.set_axisbelow(True)
ax.set_facecolor("#fafafa")

# Legend styling
handles, labels = ax.get_legend_handles_labels()
ax.legend(
    handles,
    labels,
    title="Protein Quant.",
    fontsize=12,
    title_fontsize=13,
    loc='center left',
    bbox_to_anchor=(1.01, 0.5),
    frameon=False
)

# Remove top/right spines for a cleaner look
sns.despine(ax=ax, top=True, right=True, left=False, bottom=False)

plt.tight_layout()
plots.finalize_plot(
    fig, show=True, save=save_to_folder, 
    filename='PCA_CumulativeVariance_Proteins',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)